# coding=utf-8

import struct
import json
import numpy
from numpy import array
from numpy import asarray
from numpy import concatenate
from numpy import floor
from numpy import fabs
from numpy import fromfile
from numpy import reshape
from numpy import float64
from numpy import arange
from numpy import add
from numpy import mean
from numpy import zeros
from numpy import shape
from numpy import sum
from numpy import sqrt
from numpy import std

import matplotlib.pyplot as plt


class PFD(object):
    """
    private variables:
        candidateFileName
        长度
        numdms: number of dms dm长度
        numperiods: number of periods periods长度
        numpdots: number of pdots pdots长度
        nsub: number of frequency subbands folded 折叠频域长度
        npart: number of folds in time over integration 折叠时间长度
        proflen: number of bins per profile 一行，也就是一个profile中bin数量，一行一般就是两个相位
        numchan: number of channels for radio data
        最小步长
        pstep: minimum period stepsize in profile phase bins
        pdstep: minimum p-dot stepsize in profile phase bins
        dmstep: minimum dm stepsize in profile phase bins
        ndmfact:
        npfact
        filenm: filename of the folded data
        candnm: string describing the candidate
        telescope: telescope where observation took place
        pgdev: PGPLOT device to use
        赤纬赤经等空间位置信息
        rastr: J2000 RA string in format hh:mm:ss.ssss
        decstr: J20000 DEC string in format dd:mm::ss.ssss
        dt: sampling interval of the data
        startT: fraction of observation file to start folding
        endT: fraction of observation file to stop folding
        tepoch: topocentric eopch of data in MJD
        bepoch: barycentric eopch of data in MJD
        avgvoverc: average topocentric velocity
        lofreq: center of low frequency radio channel
        chan_wid: width of each radio channel in MHZ
        bestdm: best DM
        topo_pow
        topo_p1
        topo_p2
        topo_p3
        bary_pow
        bary_p1
        bary_p2
        bary_p3
        fold_pow
        fold_p1
        fold_p2
        fold_p3
        orb_p
        orb_e
        orb_x
        orb_w
        orb_t
        orb_pd
        orb_wd
        dms
        periods
        pdots
        profs
        stats
        pts_per_fold
        subbands:the sub-band data
    """

    def __init__(self, _candidateFileName):
        self.candidateFileName = _candidateFileName
        candidateFile = open(self.candidateFileName, "rb")

        swapChar = '<'
        data = candidateFile.read(5 * 4)
        testSwap = struct.unpack(swapChar + "i" * 5, data)
        if (fabs(asarray(testSwap))).max() > 100000:
            swapChar = '>'

        (self.numdms, self.numperiods, self.numpdots, self.nsub, self.npart) = struct.unpack(swapChar + "i" * 5, data)
        (self.proflen, self.numchan, self.pstep, self.pdstep, self.dmstep, self.ndmfact, self.npfact) = struct.unpack(
            swapChar + "i" * 7, candidateFile.read(7 * 4))

        self.filenm = candidateFile.read(struct.unpack(swapChar + "i", candidateFile.read(4))[0])
        self.candnm = candidateFile.read(struct.unpack(swapChar + "i", candidateFile.read(4))[0])
        self.telescope = candidateFile.read(struct.unpack(swapChar + "i", candidateFile.read(4))[0])
        self.pgdev = candidateFile.read(struct.unpack(swapChar + "i", candidateFile.read(4))[0])

        test = candidateFile.read(16)
        has_posn = 1
        for ii in range(16):
            if str(test[ii]) not in '0123456789:.-\0':
                has_posn = 0
                break

        if has_posn:
            self.rastr = test[:test.find('\0')]
            test = candidateFile.read(16)
            self.decstr = test[:test.find('\0')]
            (self.dt, self.startT) = struct.unpack(swapChar + "dd", candidateFile.read(2 * 8))
        else:
            self.rastr = "Unknown"
            self.decstr = "Unknown"
            (self.dt, self.startT) = struct.unpack(swapChar + "dd", test)

        (self.endT, self.tepoch, self.bepoch, self.avgvoverc, self.lofreq, self.chan_wid, self.bestdm) = struct.unpack(
            swapChar + "d" * 7, candidateFile.read(7 * 8))
        (self.topo_pow, tmp) = struct.unpack(swapChar + "f" * 2, candidateFile.read(2 * 4))
        (self.topo_p1, self.topo_p2, self.topo_p3) = struct.unpack(swapChar + "d" * 3, candidateFile.read(3 * 8))
        (self.bary_pow, tmp) = struct.unpack(swapChar + "f" * 2, candidateFile.read(2 * 4))
        (self.bary_p1, self.bary_p2, self.bary_p3) = struct.unpack(swapChar + "d" * 3, candidateFile.read(3 * 8))
        (self.fold_pow, tmp) = struct.unpack(swapChar + "f" * 2, candidateFile.read(2 * 4))
        (self.fold_p1, self.fold_p2, self.fold_p3) = struct.unpack(swapChar + "d" * 3, candidateFile.read(3 * 8))
        (self.orb_p, self.orb_e, self.orb_x, self.orb_w, self.orb_t, self.orb_pd, self.orb_wd) = struct.unpack(
            swapChar + "d" * 7, candidateFile.read(7 * 8))

        self.dms = asarray(struct.unpack(swapChar + "d" * self.numdms, candidateFile.read(self.numdms * 8)))
        self.periods = asarray(struct.unpack(swapChar + "d" * self.numperiods, candidateFile.read(self.numperiods * 8)))
        self.pdots = asarray(struct.unpack(swapChar + "d" * self.numpdots, candidateFile.read(self.numpdots * 8)))

        if (swapChar == '<'):  # little endian
            self.profs = zeros((self.npart, self.nsub, self.proflen), dtype='d')
            for ii in range(self.npart):
                for jj in range(self.nsub):
                    try:
                        self.profs[ii, jj, :] = fromfile(candidateFile, float64, self.proflen)
                    except Exception:  # Catch *all* exceptions.
                        pass
                        # print ""
        else:
            self.profs = asarray(struct.unpack(swapChar + "d" * self.numprofs * self.proflen,
                                               candidateFile.read(self.numprofs * self.proflen * 8)))
            self.profs = reshape(self.profs, (self.npart, self.nsub, self.proflen))

        self.stats = zeros((self.npart, self.nsub, 7), dtype='d')

        self.pts_per_fold = []
        for ii in range(self.npart):
            currentstats = self.stats[ii]
            for jj in range(self.nsub):
                if (swapChar == '<'):  # little endian
                    try:
                        currentstats[jj] = fromfile(candidateFile, float64, 7)
                    except Exception:  # Catch *all* exceptions.
                        pass
                        # print ""
                else:
                    try:
                        currentstats[jj] = asarray(struct.unpack(swapChar + "d" * 7, candidateFile.read(7 * 8)))
                    except Exception:  # Catch *all* exceptions.
                        pass
                        # print ""
            self.pts_per_fold.append(self.stats[ii][0][0])  # numdata from foldstats

        candidateFile.close()

        # Calculate some other features
        # sub bands
        self.subbands = self.profs.sum(0)
        # sub integrations
        self.subints = self.profs.sum(1)
        # pulse profile
        self.pulse_profile = self.profs.sum(0).sum(0)




# test code
def main():
    numpy.set_printoptions(threshold=1000000000)
    testPFDFile = PFD("E:\\LOTAAS_1_PFD_Candidate_1.pfd")
    # 样例数据中dm period pdot图标的长度都是201
    print(testPFDFile.numdms, testPFDFile.numperiods, testPFDFile.numpdots)

    # bin数量100 band数量32 integration数量30
    print(testPFDFile.proflen, testPFDFile.nsub, testPFDFile.npart)

    # dms,periods和pdots都是201维数组
    print(testPFDFile.periods.shape)
    print(testPFDFile.pdots.shape)
    print(testPFDFile.dms.shape)

    # 三维原始观测数据 分别是时域 频域 相位
    # 只需要累加一个维度数据，就可得到展示在ps文件中的二维数据
    # 例如：累加时域，就是subbands图；累加频域，就是subintegration图
    # 累加时域和频域，就是脉冲特征
    print(testPFDFile.profs.shape)

    subbands = testPFDFile.subbands
    subints = testPFDFile.subints
    print(subbands.shape, subints.shape)

    pulse_profile = testPFDFile.pulse_profile
    print(pulse_profile.shape)
    print(subints)
    plt.imshow(subints)

    # plt.imshow(subints)
    # plt.plot(range(len(pulse_profile)), pulse_profile)
    plt.show()

if __name__ == "__main__":
    main()